% 二维定态Schrodinger方程的本征值法
% 二维情况下，波函数的简并态个数增加，因此不一定能猜到正确的激发态解
% 参考：https://zhuanlan.zhihu.com/p/393374195
% 可能有bug
% Gitee Repo

clc
clear

global n

L = 3;
dx = 0.05;
[x,y] = meshgrid(-L:dx:L);
n = size(x,1);

function lap = compute_laplacian()
    global n

    lap = spalloc(n*n,n*n);

    [j,i] = meshgrid(2:n-1);
    i = i(:);j = j(:);
    ind = zeros((n-2)^2,1);
    ind(:,1) = sub2ind([n,n],i,j);
    ind(:,2) = sub2ind([n,n],i+1,j);
    ind(:,3) = sub2ind([n,n],i-1,j);
    ind(:,4) = sub2ind([n,n],i,j+1);
    ind(:,5) = sub2ind([n,n],i,j-1);

    ind2 = sub2ind([n^2,n^2],ind(:,1),ind(:,1));
    lap(ind2)= -4;

    for k = 2:5
        ind2 = sub2ind([n^2,n^2],ind(:,1),ind(:,k));
        lap(ind2) = 1;
    end

    i = (1:n)';
    k0 = 1+0*i;
    k1 = n+0*i;

    ind = sub2ind([n,n],k0,i);
    ind2 = sub2ind([n^2,n^2],ind,ind);
    lap(ind2) = 1;

    ind = sub2ind([n,n],k1,i);
    ind2 = sub2ind([n^2,n^2],ind,ind);
    lap(ind2) = 1;

    ind = sub2ind([n,n],i,k0);
    ind2 = sub2ind([n^2,n^2],ind,ind);
    lap(ind2) = 1;

    ind = sub2ind([n,n],i,k1);
    ind2 = sub2ind([n^2,n^2],ind,ind);
    lap(ind2) = 1;
end

function Vmat = compute_Vmat(V)
    global n

    Vmat = spalloc(n*n,n*n);
    [j,i] = meshgrid(2:n-1);
    i = i(:);j = j(:);
    ind = sub2ind([n,n],i,j);
    ind2 = sub2ind([n^2,n^2],ind,ind);
    Vmat(ind2) = V(ind);
    return;
end

V = (x.^2+y.^2);

lap = compute_laplacian();
Vmat = compute_Vmat(V);

hbar = 1;
m = 10;

H_hat = -hbar^2/(2*m)/(dx)^2*lap + Vmat;

[u_list,H_list] = eigs(H_hat,5,0);

figure
hold on
axis equal
i = 2;
u_norm = u_list(:,i).^2/sum(u_list(:,i).^2*dx);
u_norm = reshape(u_norm,n,n);
surf(x,y,u_norm,'EdgeColor','None')

